const { Optimizer } = require("@tensorflow/tfjs-node")
const tf = require("@tensorflow/tfjs-node")
const getData = require("./data")

const TRAIN_DIR = 'train/resources/train'
const OUTPUT_DIR = 'output'
const MOBILENET_URL = 'http://ai-sample.oss-cn-hangzhou.aliyuncs.com/pipcook/models/mobilenet/web_model/model.json'
const main = async () =>{
  // 加载训练资源
  const {ds, classes} = await getData(TRAIN_DIR, OUTPUT_DIR)
  // console.log(xs, ys, classes)
  
  // 定义模型 截断模型 + 双层神经网络
  const mobilenet = await tf.loadLayersModel(MOBILENET_URL)
  // mobilenet.summary()
  const model = tf.sequential();
  for (let i = 0; i <= 86; i++) {
    const layer = mobilenet.layers[i];
    layer.trainable = false
    model.add(layer)
  }
  //
  model.add(tf.layers.flatten())
  // 隐藏层
  model.add(tf.layers.dense({
    // 神经元个数
    units: 10,
    // 激活函数
    activation: 'relu'
  }))
  // 输出层
  model.add(tf.layers.dense({
    // 神经元个数
    units: classes.length,
    // 激活函数
    activation: 'softmax'
  }))

  // 训练模型
  model.compile({
    loss: 'sparseCategoricalCrossentropy',
    optimizer: tf.train.adam(),
    metrics: ['acc']
  })
  await model.fitDataset(ds, { epochs: 20})
  await model.save(`file://${process.cwd()}/${OUTPUT_DIR}`)
}
main()